import os
import time
import numpy as np
from sklearn.preprocessing import normalize
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import torch
import torchvision as tv
import torchvision.transforms as transforms
from args import *
from tool import *
import evaluate
# from models import *


save_path = os.path.join(args.log_path, "image.{}".format(args.model_type))
if not os.path.exists(save_path):
    os.makedirs(save_path)


test_loader = torch.utils.data.DataLoader(
    tv.datasets.MNIST(args.data_path, train=False,
                   transform=transforms.Compose([
                    #    transforms.Resize([224, 224]),
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,)),
                    #    transforms.Lambda(lambda x: x.repeat(3, 1, 1))
                   ])),
    batch_size=100, shuffle=True)


P_list = (
    "a-softmax_1.0m",
    "a-softmax_2.0m",
    "a-softmax_3.0m",
    "a-softmax_0bias_0.0lambda_20scale_1dy-scale"
)
P_list = [os.path.join(args.log_path, _p) for _p in P_list]
# N_PATH = len(P_list)

M_list = (
    "19e_0.9901acc.pth",
    "15e_0.992acc.pth",
    "10e_0.9924acc.pth",
    "23e_0.9917acc.pth"
)
M_list = [os.path.join(_p, _m) for _p, _m in zip(P_list, M_list)]
M_list = [torch.load(_m).cuda() for _m in M_list]

C_list = ["C.{}e.npy".format(i) for i in (19, 15, 10, 23)]
C_list = [os.path.join(_p, _c) for _p, _c in zip(P_list, C_list)]
C_list = [np.load(_c) for _c in C_list]
N_CLASS = C_list[0].shape[0]


def get_fea():
    Fs_list, L = [], []
    for i in range(len(M_list)):
        Fs_list.append([])
    _cnt = 0
    with torch.no_grad():
        for X, Y in test_loader:
            # X, Y = X.cuda(), Y.cuda()
            X = X.cuda()
            for m, f_list in zip(M_list, Fs_list):
                _f, _ = m(X)
                f_list.append(_f.cpu().numpy())
            L.append(Y.numpy())
            _cnt += _f.size(0)
            if _cnt >= 2000:
                break

    for i in range(len(Fs_list)):
        Fs_list[i] = np.vstack(Fs_list[i])
    L = np.concatenate(L)
    return Fs_list, L


Fs_list, L = get_fea()
W_list = []  # [10, 2] x 4
for model in M_list:
    W_list.append(model.state_dict()["clf_layer.weight"].cpu().numpy())
title_list = ("margin = 1", "margin = 2", "margin = 3", "margin = 4")


W = M_list[-1].state_dict()["clf_layer.weight"].cpu().numpy()  # [10, 2]


title = "feature and centre"
fig, ax = plt.subplots(2, 2)
plt.title(title)
ax_list = []
for i in range(2):
    for j in range(2):
        ax_list.append(ax[i][j])
for _ax, _t, _F, _C in zip(ax_list, title_list, Fs_list, C_list):
    _ax.set_title(_t)
    _ax.scatter(_F[:, 0], _F[:, 1],
        s=10, c=L, marker='.', cmap="rainbow")
    for c in _C:
        _ax.plot([0, c[0]], [0, c[1]])
fig.savefig(os.path.join(save_path, "{}.png".format(title.replace(" ", "_"))))
plt.close(fig)


title = "feature and weight vector"
fig, ax = plt.subplots(2, 2)
plt.title(title)
ax_list = []
for i in range(2):
    for j in range(2):
        ax_list.append(ax[i][j])
for _ax, _t, _F, _W in zip(ax_list, title_list, Fs_list, W_list):
    _ax.set_title(_t)
    _ax.scatter(_F[:, 0], _F[:, 1],
        s=10, c=L, marker='.', cmap="rainbow")
    for w in _W:
        _ax.plot([0, w[0]], [0, w[1]])
fig.savefig(os.path.join(save_path, "{}.png".format(title.replace(" ", "_"))))
plt.close(fig)
